#ifndef _ALGO_EKF_H
#define _ALGO_EKF_H

#include "algo_matrix.h"


// #define EKF_UPDATE_CONTINUOUS
#define EKF_UPDATE_DISCRETE


// Assume the state dimension is n
// and the measurement dimension is m
#define STATE_DIM   2    // MATRIX_ROW
#define MEASURE_DIM 2    // MATRIX_COLUMN

// EKF structure definition
struct ekf_filter {
    int           state_dim;      // State dimension
    int           measure_dim;    // Measurement dimension
    Matrix_fp16_t X;              // State vector
    Matrix_fp16_t P;              // Error covariance matrix
    Matrix_fp16_t Q;              // Process noise covariance matrix
    Matrix_fp16_t R;              // Measurement noise covariance matrix
    Matrix_fp16_t F;              // State transition matrix
    Matrix_fp16_t H;              // Measurement matrix
    Matrix_fp16_t K;              // Kalman gain
    // The following are temporary matrices used to store intermediate results during
    // calculations
    Matrix_fp16_t Xdot;    // State derivative
    Matrix_fp16_t Pdot;    // Error covariance derivative
    Matrix_fp16_t E;       // Intermediate matrix for measurement update
    Matrix_fp16_t err;     // Intermediate matrix for error calculation
};


/**
 * @brief State transition function for the EKF.
 *
 * This function defines the state transition model for the system.
 *
 * @param u Input control vector.
 * @param X Current state vector.
 * @param dt Time step.
 * @param Xdot Output state derivative vector.
 * @param F Output state transition matrix.
 */
typedef void(code *ekf_sfun_t)(struct ekf_filter *kf, const float *u, const float dt);
/**
 * @brief Measurement function for the EKF.
 *
 * This function defines the measurement model for the system.
 *
 * @param y Measurement vector.
 * @param err Output measurement error vector.
 * @param X Current state vector.
 * @param H Output measurement matrix.
 */
typedef void(code *ekf_mfun_t)(struct ekf_filter *kf, const Matrix_fp16_t *y);


/**
 * @brief Initializes the Extended Kalman Filter with initial parameters.
 *
 * This function sets up the initial state, error covariance, process noise covariance,
 * and measurement noise covariance matrices. It also initializes the state transition
 * and measurement matrices to zero, which should be updated during the filter process.
 *
 * @param kf Pointer to the EKF filter structure to be initialized.
 * @param initial_state Pointer to the initial state vector.
 * @param initial_P Pointer to the initial error covariance matrix.
 * @param Q Pointer to the process noise covariance matrix.
 * @param R Pointer to the measurement noise covariance matrix.
 */
void ekf_init(struct ekf_filter *kf,
              const float       *initial_state,
              const float       *initial_P,
              const float       *Q,
              const float       *R);
/**
 * @brief Predicts the state and error covariance using the EKF model.
 *
 * This function applies the EKF prediction step, updating the state vector and
 * error covariance matrix based on the given control input and the time step.
 * It uses the state transition function defined by the user.
 *
 * @param kf Pointer to the EKF filter structure.
 * @param u Pointer to the control input vector.
 */
void ekf_filter_predict(struct ekf_filter *kf, float *u);

/**
 * @brief Updates the state and error covariance using the EKF model.
 *
 * This function applies the EKF update step, incorporating the new measurement
 * into the state estimate. It updates the state vector, error covariance matrix,
 * and computes the Kalman gain.
 *
 * @param kf Pointer to the EKF filter structure.
 * @param y Pointer to the measurement vector.
 */
void ekf_filter_update(struct ekf_filter *kf, Matrix_fp16_t *y);

void ekf_test_dynamic_filtering(void);

#endif /* EKF_H */
